Add affine AutoEP checkpoint placement - #8544

Follow up #8385

This pull request introduces support for flexible expert placement in DeepSpeed's AutoEP (Automatic Expert Placement) system, enabling non-uniform, non-contiguous, and replicated expert layouts. The main changes add a versioned expert placement descriptor, validation logic, and integration into checkpoint consolidation and metadata validation. This lays the groundwork for more advanced expert scheduling and model parallelism strategies.

The most important changes are:

AutoEP Expert Placement Descriptor and Affine Map Lowering

• Added a new module autoep_affine.py that defines the expert placement descriptor, validation, legacy uniform descriptor synthesis, and lowering to affine maps for sharded tensor reconstruction. This enables flexible, versioned expert placement beyond the legacy uniform contiguous layout.

Integration into Checkpoint Consolidation and Metadata

• Updated autoep_universal.py to:
  • Accept and validate the new expert placement descriptor in layer metadata, relaxing the requirement that num_local_experts * ep_size == num_experts when a placement is provided.
  • Use the placement descriptor and affine map for reconstructing full expert tensors during checkpoint consolidation, supporting arbitrary expert layouts. [1] [2]
  • Validate placement consistency during expert file consolidation.
  • Import and use the new placement logic. [1] [2]

Metadata Validation Enhancements

• Updated autoep_zero3_metadata.py to:
  • Import and use placement validation and legacy descriptor synthesis.
  • Track and validate placements for each layer entry, and check runtime layer consistency. [1] [2] [3]
  • Handle the presence or absence of the placement descriptor during partitioned metadata validation.

Documentation Updates

• Expanded the affine IR specification (affine_ir_spec.md) to document the new AutoEP placement descriptor, its semantics, and its integration into the IR and runtime, clarifying the distinction between placement provenance and scheduling.

Bugfixes and Robustness

• Fixed a potential bug in affine.py by skipping empty piece lists during tensor rebuilding, preventing errors when a rank has no assigned pieces.

Related: #8252, #8230.